import open3d as o3d
import numpy as np
import laspy

# 1. 读取 LAS 文件
las_file = laspy.read("/mnt/d/temp_files/part_pointCloud.las")

# 2. 提取点云数据
points = np.vstack((las_file.x, las_file.y, las_file.z)).transpose()

# 3. 创建一个 Open3D 点云对象
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)

# 4. 提取语义标签 (假设它保存在LAS的 'classification' 字段中)
# 注意：LAS标准中这个字段通常是整数类型
semantic_labels = np.asarray(las_file.classification)

# 5. 根据语义标签为点云着色
# 我们需要定义一个颜色映射函数，将标签值（0-7）映射到RGB颜色
def map_label_to_color(labels):
    """
    定义一个简单的颜色映射。
    DALES 官方类别: 0-地面, 1-植被, 2-汽车, 3-卡车, 4-电线, 5-电线杆, 6-立面, 7-杂波
    你可以根据喜好修改这些颜色！
    """
    # 创建一个空的颜色数组 [R, G, B]
    colors = np.zeros((labels.shape[0], 3))
    
    # 定义颜色映射 (这里只是示例，请调整为你喜欢的颜色)
    color_map = {
        0: [0.5, 0.5, 0.5],  # 地面 - 灰色
        1: [0.0, 1.0, 0.0],  # 植被 - 绿色
        2: [1.0, 0.0, 0.0],  # 汽车 - 红色
        3: [1.0, 1.0, 0.0],  # 卡车 - 黄色
        4: [0.0, 0.0, 1.0],  # 电线 - 蓝色
        5: [0.5, 0.0, 0.5],  # 电线杆 - 紫色
        6: [1.0, 0.5, 0.0],  # 立面 - 橙色
        7: [0.0, 0.0, 0.0],  # 杂波 - 黑色
    }
    
    for label, color in color_map.items():
        # 为每个对应标签的点分配颜色
        colors[labels == label] = color
        
    return colors

# 为点云赋色
pcd.colors = o3d.utility.Vector3dVector(map_label_to_color(semantic_labels))

# 6. 可视化！
print("正在显示点云... 窗口出来后，你可以：")
print(" - 拖动鼠标: 旋转视角")
print(" - 滚轮: 缩放")
print(" - 鼠标左键 + Ctrl: 平移")
print(" - 窗口右上角的『X』: 关闭窗口并继续执行程序")

o3d.visualization.draw_geometries([pcd],
                                  window_name="DALES Dataset Visualization",
                                  width=1024,
                                  height=768)